{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LlamaIndex\n",
    "\n",
    "## Overview\n",
    "\n",
    "This is a Quick Start guide that shows how to use Guardrails alongside LlamaIndex.  As you'll see, the LlamaIndex portion comes directly from their starter examples [here](https://docs.llamaindex.ai/en/stable/getting_started/starter_example/).  Our approach to intergration for LlamaIndex, similar to our LangChain integration, is the make the interaction feel as native to the tool as possible."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installation\n",
    "Install LlamaIndex and a version of Guardrails with LlamaIndex support."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install llama-index -q\n",
    "! pip install \"guardrails-ai>=0.6.1\" -q"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install a couple validators from the Guardrails Hub that we'll use to guard the query outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Installing hub:\u001b[35m/\u001b[0m\u001b[35m/guardrails/\u001b[0m\u001b[95mdetect_pii...\u001b[0m\n",
      "✅Successfully installed guardrails/detect_pii version \u001b[1;36m0.0\u001b[0m.\u001b[1;36m5\u001b[0m!\n",
      "\n",
      "\n",
      "Installing hub:\u001b[35m/\u001b[0m\u001b[35m/guardrails/\u001b[0m\u001b[95mcompetitor_check...\u001b[0m\n",
      "✅Successfully installed guardrails/competitor_check!\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "! guardrails hub install hub://guardrails/detect_pii hub://guardrails/competitor_check --no-install-local-models -q"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download some sample data from the LlamaIndex docs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
      "                                 Dload  Upload   Total   Spent    Left  Speed\n",
      "100 75042  100 75042    0     0   353k      0 --:--:-- --:--:-- --:--:--  354k\n"
     ]
    }
   ],
   "source": [
    "! mkdir -p ./data\n",
    "! curl https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt > ./data/paul_graham_essay.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Index Setup\n",
    "\n",
    "First we'll load some data and build an index as shown in the [starter tutorial](https://docs.llamaindex.ai/en/stable/getting_started/starter_example/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os.path\n",
    "from llama_index.core import (\n",
    "    VectorStoreIndex,\n",
    "    SimpleDirectoryReader,\n",
    "    StorageContext,\n",
    "    load_index_from_storage,\n",
    ")\n",
    "\n",
    "# check if storage already exists\n",
    "PERSIST_DIR = \"./storage\"\n",
    "if not os.path.exists(PERSIST_DIR):\n",
    "    # load the documents and create the index\n",
    "    documents = SimpleDirectoryReader(\"data\").load_data()\n",
    "    index = VectorStoreIndex.from_documents(documents)\n",
    "    # store it for later\n",
    "    index.storage_context.persist(persist_dir=PERSIST_DIR)\n",
    "else:\n",
    "    # load the existing index\n",
    "    storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)\n",
    "    index = load_index_from_storage(storage_context)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Guard Setup\n",
    "\n",
    "Next we'll create our Guard and assign some validators to check the output from our queries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from guardrails import Guard\n",
    "from guardrails.hub import CompetitorCheck, DetectPII\n",
    "\n",
    "guard = (\n",
    "    Guard()\n",
    "    .use(CompetitorCheck(competitors=[\"Fortran\", \"Ada\", \"Pascal\"], on_fail=\"fix\"))\n",
    "    .use(DetectPII(pii_entities=[\"PERSON\", \"EMAIL_ADDRESS\"], on_fail=\"fix\"))\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Querying The Index\n",
    "\n",
    "To demonstrate it's plug-and-play capabilities, first we'll query the index un-guarded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The author is Paul Graham. Growing up, he worked on writing short stories and programming, starting with the IBM 1401 in 9th grade using an early version of Fortran. Later, he transitioned to microcomputers like the TRS-80 and began programming more extensively, creating simple games and a word processor.\n"
     ]
    }
   ],
   "source": [
    "# Use index on it's own\n",
    "query_engine = index.as_query_engine()\n",
    "response = query_engine.query(\"Who is the author and what did they do growing up?\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we'll set up a guarded engine, and re-query the index to see how Guardrails applies the fixes we specified when assigning our validators to the Guard."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The author is <PERSON>. Growing up, he worked on writing short stories and programming, starting with the IBM 1401 in 9th grade using an early version of [COMPETITOR]. Later, he transitioned to microcomputers like the TRS-80 and Apple II, where he wrote simple games, programs, and a word processor. \n"
     ]
    }
   ],
   "source": [
    "# Use index with Guardrails\n",
    "from guardrails.integrations.llama_index import GuardrailsQueryEngine\n",
    "\n",
    "guardrails_query_engine = GuardrailsQueryEngine(engine=query_engine, guard=guard)\n",
    "\n",
    "response = guardrails_query_engine.query(\n",
    "    \"Who is the author and what did they do growing up?\"\n",
    ")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The GuardrailsEngine can also be used with LlamaIndex's chat engine, not just the query engine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The author is <PERSON>. Growing up, he worked on writing short stories and programming. He started with early attempts on an IBM 1401 using [COMPETITOR] in 9th grade. Later, he transitioned to microcomputers, building a Heathkit kit and eventually getting a TRS-80 to write simple games and programs. Despite enjoying programming, he initially planned to study philosophy in college but eventually switched to AI due to a lack of interest in philosophy courses. \n"
     ]
    }
   ],
   "source": [
    "# For chat engine\n",
    "from guardrails.integrations.llama_index import GuardrailsChatEngine\n",
    "\n",
    "chat_engine = index.as_chat_engine()\n",
    "guardrails_chat_engine = GuardrailsChatEngine(engine=chat_engine, guard=guard)\n",
    "\n",
    "response = guardrails_chat_engine.chat(\n",
    "    \"Tell me who the author is and what they did growing up.\"\n",
    ")\n",
    "print(response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
